package cn.kinoko.service.manga.impl;

import cn.hutool.core.util.StrUtil;
import cn.kinoko.common.code.ErrorCodeEnum;
import cn.kinoko.common.constant.Flag;
import cn.kinoko.common.utils.ExceptionUtil;
import cn.kinoko.common.utils.ImgUtil;
import cn.kinoko.common.utils.LinearRegression;
import cn.kinoko.config.minio.MinioProperties;
import cn.kinoko.model.ai.mit.MITOCRResult;
import cn.kinoko.model.ai.mit.constant.MITOCRLinearReg;
import cn.kinoko.model.ai.ocr.constant.OcrType;
import cn.kinoko.model.ai.ocr.entity.OcrLr;
import cn.kinoko.model.ai.ocr.entity.OcrLrCalibrate;
import cn.kinoko.model.manga.constant.EpCutoutStatus;
import cn.kinoko.model.manga.constant.EpMarkStatus;
import cn.kinoko.model.manga.entity.MangaEpRaw;
import cn.kinoko.model.manga.entity.MangaMark;
import cn.kinoko.service.ai.mit.MITService;
import cn.kinoko.service.ai.ocr.OcrLrCalibrateService;
import cn.kinoko.service.ai.ocr.OcrLrService;
import cn.kinoko.service.manga.MangaAIService;
import cn.kinoko.service.manga.MangaEpRawService;
import cn.kinoko.service.manga.MangaMarkService;
import cn.kinoko.service.minio.MinioService;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @author kinoko
 */
@Slf4j
@Service
@AllArgsConstructor
public class MangaAIServiceImpl implements MangaAIService {

    private MITService mitService;
    private MangaMarkService mangaMarkService;
    private MangaEpRawService mangaEpRawService;
    private OcrLrCalibrateService ocrLrCalibrateService;
    private OcrLrService ocrLrService;
    private MinioService minioService;
    private MinioProperties minioProperties;

    /**
     * 是否可用
     * @return true-可用 false-不可用
     */
    @Override
    public boolean available() {
        return mitService.available();
    }

    /**
     * 一键标注图片
     * @param pageId 页码id
     * @return 执行结果
     */
    @Transactional(rollbackFor = Exception.class)
    @Override
    public boolean ocrMark(Long pageId) {
        // 查询页码
        MangaEpRaw epRaw = mangaEpRawService.getById(pageId);
        ExceptionUtil.isNull(epRaw, ErrorCodeEnum.SYSTEM_ERROR_B0402);
        return ocrMark(epRaw);
    }

    /**
     * 一键标注图片
     * @param epRaw 页码
     * @return 执行结果
     */
    @Transactional(rollbackFor = Exception.class)
    @Override
    public boolean ocrMark(MangaEpRaw epRaw) {
        // 翻译图片
        List<MITOCRResult> results = mitService.ocrTranslate(epRaw.getImageUrl());
        // 查询线性回归系数
        OcrLr ocrLr = ocrLrService.getById(OcrType.MIT);
        ExceptionUtil.isNull(ocrLr, ErrorCodeEnum.SYSTEM_ERROR_B0409);
        // 将页码翻译状态置为已翻译
        mangaEpRawService.lambdaUpdate().eq(MangaEpRaw::getEpRid, epRaw.getEpRid())
                .set(MangaEpRaw::getMarkStatus, EpMarkStatus.DONE.getCode())
                .update();
        // 转换结果
        Map<Integer, MangaMark> newMarkMap = new HashMap<>(results.size());
        for (int index = 0; index < results.size(); index++) {
            MITOCRResult result = results.get(index);
            // 构建标注，加入集合
            MangaMark mangaMark = buildMangaMark(result, index, epRaw, ocrLr);
            newMarkMap.put(index, mangaMark);
        }
        mangaMarkService.saveBatch(newMarkMap.values());
        // 组装线性回归数据
        List<OcrLrCalibrate> ocrLrCalibrates = new ArrayList<>(results.size());
        for (int index = 0; index < results.size(); index++) {
            MangaMark mangaMark = newMarkMap.get(index);
            MITOCRResult result = results.get(index);
            OcrLrCalibrate lrCalibrate = new OcrLrCalibrate();
            lrCalibrate.setOcrType(OcrType.MIT);
            lrCalibrate.setMarkId(mangaMark.getMarkId());
            lrCalibrate.setTX(mangaMark.getX());
            lrCalibrate.setTY(mangaMark.getY());
            lrCalibrate.setOcrX(result.getMinX());
            lrCalibrate.setOcrY(result.getMinY());
            ocrLrCalibrates.add(lrCalibrate);
        }
        return ocrLrCalibrateService.saveBatch(ocrLrCalibrates);
    }

    /**
     * 确认OCR校对
     * @param pageId 页码id
     * @return 执行结果
     */
    @Transactional(rollbackFor = Exception.class)
    @Override
    public boolean ocrCheck(Long pageId) {
        // 查询页码标注
        Map<Long, MangaMark> markMap = mangaMarkService.lambdaQuery().eq(MangaMark::getEpRid, pageId).list()
                .stream().collect(Collectors.toMap(MangaMark::getMarkId, m -> m));
        // 查询ocr识图坐标
        List<OcrLrCalibrate> ocrLrCalibrates = ocrLrCalibrateService.lambdaQuery()
                .in(OcrLrCalibrate::getMarkId, markMap.keySet())
                .list();
        // 更新校准后的坐标
        for (OcrLrCalibrate ocrLrCalibrate : ocrLrCalibrates) {
            MangaMark mangaMark = markMap.get(ocrLrCalibrate.getMarkId());
            if (mangaMark == null) continue;
            ocrLrCalibrate.setTX(mangaMark.getX());
            ocrLrCalibrate.setTY(mangaMark.getY());
            ocrLrCalibrate.setStatus(Flag.YES);
        }
        return ocrLrCalibrateService.updateBatchById(ocrLrCalibrates);
    }

    /**
     * 抠除文本
     *
     * @param epRaw 页码
     */
    @Transactional(rollbackFor = Exception.class)
    @Override
    public void cutoutText(MangaEpRaw epRaw) throws IOException {
        // 抠图
        String path = mitService.cutoutText(epRaw.getImageUrl());
        // 下载资源上传至minio
        try (InputStream is = ImgUtil.download(path, 1)) {
            String key = minioProperties.path2Key(epRaw.getImageUrl());
            // 存储地址为 blank/ + 原文件路径
            String blankUrl = minioService.uploadFile(MinioProperties.BLANK_PREFIX + StrUtil.subBefore(key, "/", true),
                    "blank_" + StrUtil.subAfter(key, "/", true),
                    "image/jpeg", is,
                    false,
                    false,
                    false);
            // 将文件引用置为引用
            minioService.saveFileReference(blankUrl);
            // 将页码置为已抠除
            epRaw.setCutoutStatus(EpCutoutStatus.DONE.getCode());
            // 设置空白图片地址
            epRaw.setBlankUrl(blankUrl);
            mangaEpRawService.updateById(epRaw);
        }
    }

    /**
     * 构建标注
     * @param result 结果
     * @param index  索引
     * @param epRaw  页码
     * @param ocrLr  线性回归系数
     * @return 结果
     */
    private @NotNull MangaMark buildMangaMark(MITOCRResult result, int index, MangaEpRaw epRaw, OcrLr ocrLr) {
        MangaMark mangaMark = new MangaMark();
        mangaMark.setMangaId(epRaw.getMangaId());
        mangaMark.setEpId(epRaw.getEpId());
        mangaMark.setEpRid(epRaw.getEpRid());
        mangaMark.setText(result.getCnText());
        mangaMark.setX(Math.min((int) LinearRegression.predict(result.getMinX(), ocrLr.getXSlopeDouble(), ocrLr.getXInterceptDouble()), MITOCRLinearReg.MAX_X));
        mangaMark.setY(Math.min((int) LinearRegression.predict(result.getMinY(), ocrLr.getYSlopeDouble(), ocrLr.getYInterceptDouble()), MITOCRLinearReg.MAX_Y));
        mangaMark.setNumber(index + 1);
        return mangaMark;
    }
}